{
 "nbformat": 4,
 "nbformat_minor": 0,
 "metadata": {
  "colab": {
   "provenance": [],
   "authorship_tag": "ABX9TyP9+BJLzaiLp67cp7+DjUBb",
   "include_colab_link": true
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "view-in-github",
    "colab_type": "text"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/langroid/langroid/blob/main/examples/Langroid_QuickStart_OpenAI_Assistants_API.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Multi-Agent programming with Langroid, using the new OpenAI Assistant API\n",
    "\n",
    "OpenAI's [Assistants API](https://platform.openai.com/docs/assistants/overview) provides several conveniences to help build LLM applications, such as:\n",
    "- managing conversation state (threads)\n",
    "- persistent threads and assistants\n",
    "- tools (function-calling, retrieval, code-interpreter)\n",
    "\n",
    "\n",
    "There is a new programming paradigm emerging, where these assistants are primitives, and a key chalenge is:\n",
    "\n",
    "> how can you have these assistants collaborate to solve a task?\n",
    "\n",
    "[Langroid](https://github.com/langroid/langroid)'s new `OpenAIAssistant` class offers this ability. Langroid was designed from the start to support a multi-agent LLM programming paradigm, where agents can collaborate on a task via conversation.\n",
    "The new `OpenAIAssistant` agent gives you:\n",
    "\n",
    "- 1️⃣ a dead-simple interface to the Assistants API,\n",
    "- 2️⃣ a seamless way to have assistants collaborate with each other or with users.\n",
    "\n",
    "The Assistant API fits naturally into Langroid's notion of a `ChatAgent`,\n",
    "and the `OpenAIAssistant` class derives from `ChatAgent`.\n",
    "`OpenAIAssistant` can be used as a drop-in replacement for `ChatAgent` in any\n",
    "Langroid application, and leverage the **multi-agent** task orchestration built\n",
    "into Langroid.\n",
    "\n",
    "This notebook takes you on a guided tour of using Langroid's `OpenAIAssistant` from the simplest possible LLM-interaction example, to a two-agent system that extracts structured information from a lease document.\n",
    "\n",
    "![langroid-oai](https://github.com/langroid/langroid-examples/blob/main/examples/docqa/langroid-oai.png?raw=true)\n",
    "\n"
   ],
   "metadata": {
    "id": "b9fHPojfnbPy"
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Install, setup, import"
   ],
   "metadata": {
    "id": "psOMvEL0Gekz"
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "A8-Y_YPZutn6",
    "outputId": "2a5fb145-8ee0-4215-a442-29e75e96bdbd"
   },
   "source": [
    "# Silently install, suppress all output (~2-4 mins)\n",
    "!pip install -q --upgrade langroid &> /dev/null\n",
    "!pip show langroid"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "# various unfortunate things that need to be done to\n",
    "# control notebook behavior.\n",
    "\n",
    "# (a) output width\n",
    "\n",
    "from IPython.display import HTML, display\n",
    "\n",
    "def set_css():\n",
    "  display(HTML('''\n",
    "  <style>\n",
    "    pre {\n",
    "        white-space: pre-wrap;\n",
    "    }\n",
    "  </style>\n",
    "  '''))\n",
    "get_ipython().events.register('pre_run_cell', set_css)\n",
    "\n",
    "# (b) logging related\n",
    "import logging\n",
    "logging.basicConfig(level=logging.ERROR)\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "import logging\n",
    "for logger_name in logging.root.manager.loggerDict:\n",
    "    logger = logging.getLogger(logger_name)\n",
    "    logger.setLevel(logging.ERROR)\n",
    "\n"
   ],
   "metadata": {
    "id": "rWwH6duUzAC6"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 17
    },
    "id": "U5Jav3hPofNq",
    "outputId": "f78ffcef-6be1-4c77-d79e-0b194b297384"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### OpenAI API Key (Needs GPT4-TURBO)"
   ],
   "metadata": {
    "id": "j-6vNfKW9J7b"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "# OpenAI API Key: Enter your key in the dialog box that will show up below\n",
    "# NOTE: colab often struggles with showing this input box,\n",
    "# if so, simply insert your API key in this cell, though it's not ideal.\n",
    "import os\n",
    "\n",
    "from getpass import getpass\n",
    "\n",
    "os.environ['OPENAI_API_KEY'] = getpass('Enter your GPT4-Turbo-capable OPENAI_API_KEY key:', stream=None)\n",
    "\n",
    "\n"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "id": "uvTODlZv3yyT",
    "outputId": "3e33fdfe-d5de-46d5-e388-23bf81a04d77"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "from pydantic import BaseModel\n",
    "import json\n",
    "import os\n",
    "\n",
    "from langroid.agent.openai_assistant import (\n",
    "    OpenAIAssistantConfig,\n",
    "    OpenAIAssistant,\n",
    "    AssistantTool,\n",
    ")\n",
    "\n",
    "from langroid.agent.chat_agent import ChatAgent, ChatAgentConfig\n",
    "from langroid.agent.task import Task\n",
    "from langroid.agent.tool_message import ToolMessage\n",
    "from langroid.language_models.openai_gpt import OpenAIGPTConfig, OpenAIChatModel\n",
    "from langroid.utils.logging import setup_colored_logging\n",
    "from langroid.utils.constants import NO_ANSWER\n",
    "from langroid.utils.configuration import settings\n",
    "settings.notebook = True"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 17
    },
    "id": "A5N0NQwc3jX_",
    "outputId": "7452e570-b280-4854-a89b-c1472a8208ba"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Example 1: Basic Chat Example with Assistant API\n",
    "Langroid's `OpenAIAssistant` class helps you easily use the OpenAI Assistant API to get a response from the LLM and ask follow-up questions (note that conversation state is maintained by the Assistant API via threads).\n"
   ],
   "metadata": {
    "id": "8vDpiY0XHAkT"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "cfg = OpenAIAssistantConfig(\n",
    "    llm = OpenAIGPTConfig(chat_model=OpenAIChatModel.GPT4_TURBO)\n",
    ")\n",
    "agent = OpenAIAssistant(cfg)\n",
    "\n",
    "response = agent.llm_response(\"What is the square of 3?\")"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 172
    },
    "id": "9c5Av3rKHQIm",
    "outputId": "1ed4763f-defc-475b-e0e3-d64537b67b08"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "response = agent.llm_response(\"What about 5?\") # maintains conv state"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 33
    },
    "id": "5GvqhTlBRgXp",
    "outputId": "f4b93adb-e1c3-4a52-d1c6-f3260b94cce5"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Example 2: Wrap Agent in a Task, run it\n",
    "\n",
    "An `OpenAIAssistant` agent has various capabilities (LLM responses, agent methods/tools, etc) but there is no mechanism to iterate over these capabilities or with a human or with other agents.\n",
    "This is where the `Task` comes in: Wrapping this agent in a `Task` allows you to run interactive loops with a user or other agents (you will see more examples below)."
   ],
   "metadata": {
    "id": "-MVHyF4cSGb0"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "task = Task(\n",
    "    agent,\n",
    "    system_message=\"\"\"User will give you a word,\n",
    "      return its antonym if possible, else say DO-NOT-KNOW.\n",
    "      Be concise!\",\n",
    "      \"\"\",\n",
    "    single_round=True\n",
    ")\n",
    "result = task.run(\"ignorant\")\n"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 66
    },
    "id": "8cmc5aDzScdO",
    "outputId": "253fec3c-2f03-428b-83bc-f1170702fef0"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Example 3: OpenAIAssistant Agent + Task with Code Interpreter\n",
    "Here we attach the \"code_interpreter\" tool (from the OpenAI Assistant API) to the agent defined above, and run it in a task."
   ],
   "metadata": {
    "id": "veWSLzDSVDzB"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "agent.add_assistant_tools([AssistantTool(type=\"code_interpreter\")])\n",
    "task = Task(agent, interactive=False, single_round=True)\n",
    "result = task.run(\"What is the 10th Fibonacci number, if you start with 1,2?\")"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 428
    },
    "id": "5-h1ztagTd7Y",
    "outputId": "4bcacb3a-e9d4-4d1c-8225-c6a090711459"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Example 4: OpenAIAssistant with Retrieval\n",
    "Attach a file (a lease document) and the \"retrieval\" tool, and ask questions about the document."
   ],
   "metadata": {
    "id": "DvyNcH5HbodS"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "# get the lease document\n",
    "\n",
    "import requests\n",
    "file_url = \"https://raw.githubusercontent.com/langroid/langroid-examples/main/examples/docqa/lease.txt\"\n",
    "response = requests.get(file_url)\n",
    "with open('lease.txt', 'wb') as file:\n",
    "    file.write(response.content)\n",
    "\n",
    "# verify\n",
    "#with open('lease.txt', 'r') as file:\n",
    "#   print(file.read())\n",
    "\n",
    "# now create agent, add retrieval tool and file\n",
    "agent = OpenAIAssistant(cfg)\n",
    "agent.add_assistant_tools([AssistantTool(type=\"retrieval\")])\n",
    "agent.add_assistant_files([\"lease.txt\"])\n",
    "response = agent.llm_response(\"What is the start date of the lease?\")\n"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 172
    },
    "id": "fegAio3kpgoo",
    "outputId": "eae49d56-7f40-4480-98aa-e4e1c523a910"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Example 5: OpenAIAsssistant + Task: Custom Function-calling\n",
    "You can define your own custom function (or `ToolMessage` in Langroid terminology), enable the agent to use it, and have a special method to handle the message when the LLM emits such a message."
   ],
   "metadata": {
    "id": "Xub3BgSMc4uA"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "# Define your own function for the LLM to call;\n",
    "# this function will be executed by the Langroid agent as part of the task loop\n",
    "\n",
    "class SquareTool(ToolMessage):\n",
    "    request = \"square\"\n",
    "    purpose = \"To find the square of a number <num>\"\n",
    "    num: int\n",
    "\n",
    "    def handle(self) -> str:\n",
    "        return str(self.num ** 2)\n",
    "\n",
    "# create agent, add tool to agent\n",
    "cfg = OpenAIAssistantConfig(\n",
    "    llm=OpenAIGPTConfig(chat_model=OpenAIChatModel.GPT4_TURBO),\n",
    "    name=\"NumberExpert\",\n",
    ")\n",
    "agent = OpenAIAssistant(cfg)\n",
    "agent.enable_message(SquareTool)\n",
    "task = Task(\n",
    "    agent,\n",
    "    system_message=\"\"\"\n",
    "    User will ask you to square a number.\n",
    "    You do NOT know how, so you will use the\n",
    "    `square` function to find the answer.\n",
    "    When you get the answer say DONE and show it.\n",
    "    \"\"\",\n",
    "    interactive=False,\n",
    ")\n",
    "response = task.run(\"What is the square of 5?\")\n"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 352
    },
    "id": "dgav7-JOdAUM",
    "outputId": "b3835bfb-90ca-4642-e585-33743c5730a6"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Example 6: 2-Agent system to extract structured info from a Lease Document\n",
    "Now we are ready to put together the various notions above, to build a two-agent system where:\n",
    "- Lease Extractor Agent is required to collect structured information about a lease document, but does not have access to it, so it generates questions to:\n",
    "- Retriever Agent which answers questions it receives, using the \"retrieval\" tool, based on the attached lease document\n"
   ],
   "metadata": {
    "id": "yi9GppzlKae_"
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### Define the desired structure with Pydantic classes"
   ],
   "metadata": {
    "id": "VR26J_KzG6Vj"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "\n",
    "class LeasePeriod(BaseModel):\n",
    "    start_date: str\n",
    "    end_date: str\n",
    "\n",
    "\n",
    "class LeaseFinancials(BaseModel):\n",
    "    monthly_rent: str\n",
    "    deposit: str\n",
    "\n",
    "\n",
    "class Lease(BaseModel):\n",
    "    \"\"\"\n",
    "    Various lease terms.\n",
    "    Nested fields to make this more interesting/realistic\n",
    "    \"\"\"\n",
    "\n",
    "    period: LeasePeriod\n",
    "    financials: LeaseFinancials\n",
    "    address: str\n",
    "\n"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 17
    },
    "id": "Q6GXjhWf5DkQ",
    "outputId": "ec9c930f-245a-4151-950d-8f407b439c2c"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### Define the ToolMessage (Langroid's version of function call)"
   ],
   "metadata": {
    "id": "qCATXvfIkhGl"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "\n",
    "class LeaseMessage(ToolMessage):\n",
    "    \"\"\"Tool/function to use to present details about a commercial lease\"\"\"\n",
    "\n",
    "    request: str = \"lease_info\"\n",
    "    purpose: str = \"Collect information about a Commercial Lease.\"\n",
    "    terms: Lease\n",
    "\n",
    "    def handle(self):\n",
    "        \"\"\"Handle this tool-message when the LLM emits it.\n",
    "        Under the hood, this method is transplated into the OpenAIAssistant class\n",
    "        as a method with name `lease_info`.\n",
    "        \"\"\"\n",
    "        print(f\"DONE! Successfully extracted Lease Info:\" f\"{self.terms}\")\n",
    "        return json.dumps(self.terms.dict())"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 17
    },
    "id": "Ffi_0u-PupvO",
    "outputId": "776a2f4c-388c-4441-c618-2682a4469e37"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### Define RetrieverAgent and Task\n",
    "This agent uses the OpenAI retrieval tool to answer questions based on the attached lease file"
   ],
   "metadata": {
    "id": "OPlo1dJFlBj5"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "  retriever_cfg = OpenAIAssistantConfig(\n",
    "        name=\"LeaseRetriever\",\n",
    "        llm=OpenAIGPTConfig(chat_model=OpenAIChatModel.GPT4_TURBO),\n",
    "        system_message=\"Answer questions based on the documents provided.\",\n",
    "    )\n",
    "\n",
    "  retriever_agent = OpenAIAssistant(retriever_cfg)\n",
    "  retriever_agent.add_assistant_tools([AssistantTool(type=\"retrieval\")])\n",
    "  retriever_agent.add_assistant_files([\"lease.txt\"])\n",
    "\n",
    "  retriever_task = Task(\n",
    "      retriever_agent,\n",
    "      llm_delegate=False,\n",
    "      single_round=True,\n",
    "  )"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 156
    },
    "id": "GgzoPxX_us52",
    "outputId": "37f6d163-5980-41d8-8ecb-7e709853d5d4"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### Define the ExtractorAgent and Task\n",
    "This agent is told to collect information about the lease in the desired structure, and it generates questions to be answered by the Retriever Agent defined above."
   ],
   "metadata": {
    "id": "_m1lF9qblXj9"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "    extractor_cfg = OpenAIAssistantConfig(\n",
    "        name=\"LeaseExtractor\",\n",
    "        llm=OpenAIGPTConfig(chat_model=OpenAIChatModel.GPT4_TURBO),\n",
    "        system_message=f\"\"\"\n",
    "        You have to collect information about a Commercial Lease from a\n",
    "        lease contract which you don't have access to. You need to ask\n",
    "        questions to get this information. Ask only one or a couple questions\n",
    "        at a time!\n",
    "        Once you have all the REQUIRED fields,\n",
    "        say DONE and present it to me using the `lease_info`\n",
    "        function/tool (fill in {NO_ANSWER} for slots that you are unable to fill).\n",
    "        \"\"\",\n",
    "    )\n",
    "    extractor_agent = OpenAIAssistant(extractor_cfg)\n",
    "    extractor_agent.enable_message(LeaseMessage, include_defaults=False)\n",
    "\n",
    "    extractor_task = Task(\n",
    "        extractor_agent,\n",
    "        llm_delegate=True,\n",
    "        single_round=False,\n",
    "        interactive=False,\n",
    "    )\n",
    "\n",
    "\n",
    "\n"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 156
    },
    "id": "PV4FYnO7uxOC",
    "outputId": "e5eeed02-7785-4361-cd01-96fef92149d4"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### Add the Retriever as a subtask of Extractor, Run Extractor"
   ],
   "metadata": {
    "id": "QcA4oRaUl6oe"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "extractor_task.add_sub_task(retriever_task)\n",
    "extractor_task.run()"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "wFjUVTnCwB96",
    "outputId": "468a147b-7485-4fad-8cab-45411b18021f"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [],
   "metadata": {
    "id": "uZlas6DA0Zu6"
   },
   "execution_count": null,
   "outputs": []
  }
 ]
}
